from numpy import sin, cos
import numpy as np
from cv2 import cv2
from glob import glob
import os


class Camera:
    def __init__(self):
        if not os.path.exists("./camera_settings/cameraMatrix.txt") or not os.path.exists("./camera_settings/distCoeffs.txt"):
            print("camera files not found")
            return        
        self.cameraMatrix = np.loadtxt("./camera_settings/cameraMatrix.txt", delimiter=',')
        self.distCoeffs = np.loadtxt("./camera_settings/distCoeffs.txt", delimiter=',')
        self.h = np.loadtxt("./camera_settings/height.txt")
        self.x0 = self.cameraMatrix[0, 2]
        self.y0 = self.cameraMatrix[1, 2]
        self.fx = self.cameraMatrix[0, 0]
        self.fy = self.cameraMatrix[1, 1]
        self.f = (self.fx+self.fy)/2

    def undistort(self, uv):
        uv_undistort = cv2.undistortPoints(uv.reshape(-1, 1, 2), self.cameraMatrix, self.distCoeffs).reshape(-1,2)
        # **///需要将变换后的点先变化为齐次坐标系，然后还需要乘以相机参数矩阵，才是最终的变化的坐标。**
        uv_undistort = np.hstack([uv_undistort, np.ones((uv_undistort.shape[0], 1))])
        uv_undistort = uv_undistort.dot(self.cameraMatrix)
        uv_undistort[:, 0] += self.x0
        uv_undistort[:, 1] += self.y0
        return uv_undistort[:, 0], uv_undistort[:, 1]

    def get_dist(self, phi, omega, uv):
        u, v = self.undistort(uv)
        u_ = u - self.x0  # 像平面坐标，以像主点为原点
        v_ = self.y0 - v
        X_temp = cos(omega) * u_ + cos(phi) * sin(omega) * self.f - sin(phi) * sin(omega) * v_
        Y_temp = sin(omega) * u_ - cos(phi) * cos(omega) * self.f + sin(phi) * cos(omega) * v_
        Z_temp = sin(phi) * self.f + cos(phi) * v_
        X = X_temp / Z_temp * (-self.h)
        Y = Y_temp / Z_temp * (self.h)  # 这里我推导的是-1，但画出来结果反了，不想推了。
        return X, Y


def calib(img_path, nx, ny, cell_size):
    # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
    objp = np.zeros((ny * nx, 3), np.float32)
    objp[:, :2] = np.mgrid[0:nx, 0:ny].T.reshape(-1, 2)*cell_size
    
    # Arrays to store object points and image points from all the images.
    objpoints = []  # 3d points in real world space
    imgpoints = []  # 2d points in image plane.
    
    images = glob(img_path + '*.jpg')
    if len(images) > 0:
        print("images num for calibration : ", len(images))
    else:
        print("No image for calibration.")
        return
    
    ret_count = 0
    cv2.namedWindow("img", cv2.WINDOW_NORMAL)
    cv2.resizeWindow("img", 960, 544)
    for _, fname in enumerate(images):
        img = cv2.imread(fname)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img_size = (img.shape[1], img.shape[0])
        # Find the chessboard corners
        ret, corners = cv2.findChessboardCorners(gray, (nx, ny), None)
        cv2.drawChessboardCorners(img,(nx,ny),corners,ret)
        cv2.imshow('img', img)
        cv2.waitKey()

        # If found, add object points, image points
        if ret:
            ret_count += 1
            objpoints.append(objp)
            imgpoints.append(corners)
            print("read image:", ret_count)
    if len(objpoints) == 0:
        print("failed to find corners")
        return        
    ret, cameraMatrix, distCoeffs, _, _ = cv2.calibrateCamera(objpoints, imgpoints, img_size, None, None)
    print('Do calibration successfully')
    # 为了避免误操作覆盖掉正确的参数，先保存在图片文件夹下
    np.savetxt(img_path + "cameraMatrix.txt", cameraMatrix, fmt="%.18f", delimiter=',')
    np.savetxt(img_path + "distCoeffs.txt", distCoeffs, fmt="%.18f", delimiter=',')
    return ret, cameraMatrix, distCoeffs
    

if __name__ == "__main__":
    # 标定
    calib("./camera_calib/chessboard/", 9, 13, 192.5)

    # 校正
    cam = Camera()
    root = "./camera_calib/pic/"
    imgname = "IMG_080652_830.jpg"
    img = cv2.imread(root + imgname)
    cam.cameraMatrix = np.loadtxt("./camera_settings/cameraMatrix.txt", delimiter=',')
    cam.distCoeffs = np.loadtxt("./camera_settings/distCoeffs.txt", delimiter=',')
    img2 = cv2.undistort(img, cam.cameraMatrix, cam.distCoeffs)
    cv2.imwrite('./camera_calib/undistort/' + imgname[:-4] + "_undistort.jpg", img2)
